import sys
parent_dir = "../pygl"
sys.path.append(parent_dir)
import glfw
import OpenGL.GL as gl
import glm
import imgui
import numpy as np
from common import AppWithTrackballMenu, PointClouds, SolidWireframe, SolidLighting
import trimesh
from tkinter import filedialog
from PIL import Image

class MeshViewer(AppWithTrackballMenu):
    def setupVertex(self) -> None:
        gl.glEnable(gl.GL_DEPTH_TEST)
        self.background_color = glm.vec4(0.8, 0.8, 0.8, 1.0)
        self.opened_file = False
        self.camera.far = 800
        self.camera.zoom_speed = 0.4
        self.camera.reset(glm.vec3(2.0, 1.0, 2.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        self.imgui_window_width = 250
        self.imgui_window_height = 260
        self.point_clouds = None
        self.solid_wireframe = None
        self.solid_lighting = None
        # 数据显示方式，0 - 点云，1 - 实体线框, 2 - 实体光照
        self.view_mode = 2
        self.scale = 5.0

    def openFile(self, file_name:str)->None:
        # 用trimesh库读取.ply, .obj, .stl等格式的文件
        mesh = trimesh.load(file_name)
        self.vertices = (mesh.vertices - mesh.center_mass).astype(np.float32)
        self.indices =  mesh.faces.reshape(-1).astype(np.uint32)
        self.normals = mesh.vertex_normals.astype(np.float32)
        self.vertex_colors = None
        if isinstance(mesh.visual, trimesh.visual.ColorVisuals):
            self.vertex_colors = mesh.visual.vertex_colors
            print(self.vertex_colors[0])
        elif isinstance(mesh.visual, trimesh.visual.TextureVisuals):
            self.vertex_colors = np.full((len(self.vertices), 4), 
                                         [102, 102, 102, 255], dtype=np.uint8)
        self.point_clouds = PointClouds(self.vertices, self.vertex_colors, 2.0)
        self.point_clouds.scale = glm.vec3(self.scale)
        self.solid_wireframe = SolidWireframe(self.vertices, self.normals, 
                                              self.indices)
        self.solid_wireframe.scale = glm.vec3(self.scale)
        self.solid_lighting = SolidLighting(self.vertices, self.normals, 
                                            self.indices)
        self.solid_lighting.scale = glm.vec3(self.scale)
        self.opened_file = True

    # 保存绘图的屏幕截图到png文件
    def saveFile(self, file_name:str)->None:
        data = gl.glReadPixels(0, 0, self.width, self.height, gl.GL_RGBA, 
                               gl.GL_UNSIGNED_BYTE)
        image = Image.frombytes('RGBA', (self.width, self.height), data)
        image = image.transpose(Image.FLIP_TOP_BOTTOM)
        image.save(file_name)

    # 配置imgui界面元素
    def showImGui(self)->None:
        imgui.begin("设置")
        if imgui.button("重置相机位置"):
            self.camera.reset(glm.vec3(2.0, 1.0, 2.0), glm.vec3(0.0, 0.0, 0.0), 
                    glm.vec3(0.0, 1.0, 0.0))
        if self.opened_file:
            _, self.scale = imgui.slider_float("缩放比例", self.scale, 1.0, 20.0)
            self.point_clouds.scale = glm.vec3(self.scale)
            self.solid_wireframe.scale = glm.vec3(self.scale)
            self.solid_lighting.scale = glm.vec3(self.scale)
            if self.view_mode == 0:
                _, self.point_clouds.point_size = imgui.slider_float("点大小",
                            self.point_clouds.point_size, 1.0, 6.0)
            if self.view_mode == 1:
                _, color = imgui.color_edit3("曲面颜色", 
                            self.solid_wireframe.surface_color.x, 
                            self.solid_wireframe.surface_color.y, 
                            self.solid_wireframe.surface_color.z)
                self.solid_wireframe.surface_color = glm.vec3(color[0], 
                                                        color[1], color[2])
                _, color = imgui.color_edit3("线框颜色", 
                            self.solid_wireframe.line_color.x, 
                            self.solid_wireframe.line_color.y, 
                            self.solid_wireframe.line_color.z)
                self.solid_wireframe.line_color = glm.vec3(color[0], 
                                                           color[1], color[2])
                _, self.solid_wireframe.line_width = imgui.slider_float("线框粗细", 
                            self.solid_wireframe.line_width, 0.1, 2.0)
            if self.view_mode == 2:
                _, color = imgui.color_edit3("曲面颜色", 
                            self.solid_lighting.surface_color.x, 
                            self.solid_lighting.surface_color.y, 
                            self.solid_lighting.surface_color.z)
                self.solid_lighting.surface_color = glm.vec3(color[0], 
                                                             color[1], color[2])
                _, pos = imgui.slider_float3("光源位置", 
                            self.solid_lighting.light.position.x, 
                            self.solid_lighting.light.position.y, 
                            self.solid_lighting.light.position.z, -3.0, 3.0)
                self.solid_lighting.light.position = glm.vec3(pos)
        imgui.end()   

    # 显示imgui菜单
    def showMenu(self)->None:
        if imgui.begin_main_menu_bar():
            if imgui.begin_menu("文件"):
                clicked_open, selected_open = imgui.menu_item("打开", None, False)
                if clicked_open:
                    file_name = filedialog.askopenfilename(title="打开文件",
                        filetypes=[('ply 文件', '.ply'),
                                   ('stl 文件', '.stl'),
                                   ('obj 文件', '.obj'),
                                   ('所有文件', '*.*')]
                        )
                    if file_name != '':
                        self.openFile(file_name)
                clicked_save, selected_save = imgui.menu_item("保存", None, False)
                if clicked_save:
                    file_name = filedialog.asksaveasfilename(title="保存文件",
                        filetypes=[('png 文件', '*.png')],
                        defaultextension='png',
                        initialdir='screenshot/'
                        )
                    if file_name != '':
                        self.saveFile(file_name)
                clicked_quit, selected_quit = imgui.menu_item("退出", None, False)
                if clicked_quit:
                    glfw.set_window_should_close(self.window, True)
                imgui.end_menu()
            if imgui.begin_menu("查看方式"):
                clicked_point_clouds, selected_point_clouds = \
                        imgui.menu_item("点云", None, False)
                if clicked_point_clouds:
                    self.view_mode = 0
                clicked_solid_wireframe, selected_solid_wireframe = \
                    imgui.menu_item("实体线框", None, False)
                if clicked_solid_wireframe:
                    self.view_mode = 1
                clicked_solid_lighting, selected_solid_lighting = \
                    imgui.menu_item("实体光照", None, False)
                if clicked_solid_lighting:
                    self.view_mode = 2
                imgui.end_menu()
            imgui.end_main_menu_bar()

    def render(self) -> None:
        gl.glClearNamedFramebufferfv(0, gl.GL_COLOR, 0, 
                                     self.background_color.to_tuple())
        gl.glClearNamedFramebufferfv(0, gl.GL_DEPTH, 0, [1.0]); 
        view = self.camera.getViewMatrix()
        projection = self.camera.getProjectionMatrix()
        if self.opened_file:
            if self.view_mode == 0:
                self.point_clouds.updateModelMatrix()
                self.point_clouds.render(view, projection)
            if self.view_mode == 1:
                self.solid_wireframe.updateModelMatrix()
                self.solid_wireframe.render(view, projection)
            if self.view_mode == 2:
                self.solid_lighting.updateModelMatrix()
                self.solid_lighting.render(view, projection, self.camera.position)

    def cleanup(self) -> None:
        self.point_clouds.delete()
        self.solid_wireframe.delete()
        self.solid_lighting.delete()

if __name__ == "__main__":
    app = MeshViewer(100, 100, 1000, 600, "Mesh网格查看器")
    app.run()


